#!/usr/bin/env python3
import re
import json
import time
import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import requests

# 配置
TIMEOUT = 10          # 秒
MAX_WORKERS = 10      # 并发数
HEADERS = {
    'User-Agent': 'VLC/3.0.0'
}

def parse_m3u(file_path):
    """解析 m3u 文件，返回频道列表"""
    with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
        lines = f.readlines()

    channels = []
    extinf = None

    for line in lines:
        line = line.strip()
        if not line or line.startswith('#EXTM3U'):
            continue
        if line.startswith('#EXTINF:'):
            extinf = line
        elif line.startswith('http') and extinf:
            # 提取名称（逗号后部分）
            name_match = re.search(r',(.+)$', extinf)
            name = name_match.group(1).strip() if name_match else "Unknown"
            # 提取 group-title
            group_match = re.search(r'group-title="([^"]*)"', extinf, re.IGNORECASE)
            group = group_match.group(1) if group_match else "Other"
            channels.append({
                'extinf': extinf,
                'name': name,
                'group': group,
                'url': line
            })
            extinf = None
    return channels

def is_url_valid(url):
    """检测 URL 是否有效"""
    try:
        # 先尝试 HEAD，失败再用 GET
        response = requests.head(url, timeout=TIMEOUT, headers=HEADERS, allow_redirects=True)
        if response.status_code not in (200, 206, 301, 302):
            # 尝试 GET（有些服务器禁用 HEAD）
            response = requests.get(url, timeout=TIMEOUT, headers=HEADERS, stream=True)
        if response.status_code in (200, 206):
            content_type = response.headers.get('content-type', '').lower()
            # 允许 m3u8、视频、音频等
            if any(ct in content_type for ct in ['mpegurl', 'video', 'audio', 'octet-stream']):
                return True
    except Exception as e:
        pass
    return False

def main(input_file, output_file):
    print(f"正在解析 {input_file} ...")
    channels = parse_m3u(input_file)
    print(f"共找到 {len(channels)} 个频道")

    valid_channels = []
    total = len(channels)

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        future_to_channel = {
            executor.submit(is_url_valid, ch['url']): ch for ch in channels
        }
        for i, future in enumerate(as_completed(future_to_channel)):
            channel = future_to_channel[future]
            try:
                if future.result():
                    valid_channels.append(channel)
            except Exception as e:
                pass
            print(f"\r进度: {i+1}/{total} | 有效: {len(valid_channels)}", end='', flush=True)

    print(f"\n✅ 有效频道数: {len(valid_channels)}")

    # 写入 JSON
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(valid_channels, f, ensure_ascii=False, indent=2)

    print(f"有效频道已保存到 {output_file}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='检测 IPTV m3u 列表有效性')
    parser.add_argument('input', help='输入的 .m3u 文件路径')
    parser.add_argument('-o', '--output', default='valid_channels.json', help='输出 JSON 文件路径')
    args = parser.parse_args()

    main(args.input, args.output)